// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief Implementation of Binary BT Association (BitStrings)
*/

#include <math.h>
#include "passoc_bt.h"
#include "pmaths.h"

void passoc_bt_write_file_header(FILE *out, struct selected_options *ops) {
	if (out != NULL && ops != NULL) {	
		int ci = (int) (100 * ops->ci_level);
		fprintf(out, "CHR,BP,SNP,F_A,F_U,OR,L%i,U%i,CHISQ,P,EMP1,EMP2,BONF,HOLM,FDR_BH\n", ci, ci); 
	}
}

// Return the CHISQ value.
double passoc_bt_assoc(BOOL X, BOOL haploid, BitString ones_and_twos, struct sample *samples, int nSamples, struct small_locus *snp, FILE *out, double emp1, double emp2, double zt, struct adjusted *adj) {
	register int i;
	int A1 = 0, A2 = 0, A0 = 0;
	int U1 = 0, U2 = 0, U0 = 0;

	double chisq = 0;
	double tot, Taff, Tunf, Ta1, Ta2;
	double Texp_afffreq1, Texp_unffreq1, Texp_afffreq2, Texp_unffreq2;

	// Iterate over individuals
	for (i = 0; i < nSamples; i++) {
		unsigned int mask;
		long bitpos = (long) i * 2;
		BOOL s1 = FALSE, s2 = FALSE;

		if (samples[i].pperson->missing) continue;

		// Read S1 value
		mask = 0x1;
		mask = mask << (bitpos % BITBULKSIZE);
		mask = ones_and_twos[bitpos/BITBULKSIZE] & mask;
		if (mask !=0) s1 = TRUE;
		bitpos++;

		// Read S2 value
		mask = 0x1;
		mask = mask << (bitpos % BITBULKSIZE);
		mask = ones_and_twos[bitpos/BITBULKSIZE] & mask;
		if (mask !=0) s2 = TRUE;

		// Type of marker
		if (haploid || (X && samples[i].sex)) { 
			// Haploid marker (or male X)
			if (samples[i].pperson->aff) {	      
				if (!s1) {
					if (!s2)  A1++;
				}
				else {
					if (s2) A2++;
					else A0++; 
				}
			}
			else { // unaffected if not missing
				if (!s1) {
					if (!s2) U1++;
				} else {
					if (s2) U2++;
					else U0++; 
				}
			}
		} else {
			// Autosomal marker
			if (samples[i].pperson->aff) 
			{	      
				if (!s1) { 
					if (!s2)  A1+=2;
					else { 
						A1++; 
						A2++; 
					}  
				} else {
					if (s2) A2+=2;
					else A0+=2;  
				}
			} else { // unaffected if not missing
				if (!s1) {
					if (!s2)  U1+=2;
					else { 
						U1++; 
						U2++; 
					} 
				} else {
					if (s2) U2+=2;
					else U0+=2;  
				}
			}
		}
	}

	// Calculate the CHISQ value.
	tot = A1 + A2 + U1 + U2;
	Taff = A1+A2;
	Tunf = U1+U2;
	Ta1 = A1+U1;
	Ta2 = A2+U2;

	Texp_afffreq1 = (Taff*Ta1)/tot;
	Texp_unffreq1 = (Tunf*Ta1)/tot;
	Texp_afffreq2 = (Taff*Ta2)/tot;
	Texp_unffreq2 = (Tunf*Ta2)/tot;

	// Original not calculated so get the CHISQ value.
	chisq = 
		((A1 - Texp_afffreq1) * (A1 - Texp_afffreq1)) / Texp_afffreq1
		+ ( (U1 - Texp_unffreq1) * ( U1 - Texp_unffreq1 ) ) / Texp_unffreq1 
		+ ( (A2 - Texp_afffreq2) * ( A2 - Texp_afffreq2 ) ) / Texp_afffreq2 
		+ ( (U2 - Texp_unffreq2) * ( U2 - Texp_unffreq2 ) ) / Texp_unffreq2;

	// Capture the CHISQ for the calculation to find FDR et cetera.
	if (adj != NULL && out == NULL) if (realnum(chisq)) adj->chisq = chisq;
	
	// Write the output to the text file.
	if (out != NULL && snp != NULL && adj != NULL) {
		double sum, pvalue = 0, OR = 0.0, lowerCI = 0.0, upperCI = 0.0;
		
		pvalue = chiprobP(chisq, 1);
		fprintf(out, "%i,%i,%s,", snp->chr, snp->bp, snp->name);

		// Affected counts (I think anyway)....
		sum = A1 + A2;		
		if (A1 > 0 && sum > 0)
			fprintf(out, "%f,", ((double) A1/ sum));
		else
			fprintf(out, "NA,");

		sum = U1 + U2;		
		if (U1 > 0 && sum > 0)
			fprintf(out, "%f,", ((double) U1/ sum));
		else
			fprintf(out, "NA,");

		// Odds ratio
		if (getOddsRatio(A1, A2, U1, U2, &OR, &lowerCI, &upperCI, zt)) {
			if (lowerCI < 0) lowerCI = 0;
			if (upperCI < 0) upperCI = 0;
			fprintf(out, "%f,%f,%f,", OR, lowerCI, upperCI);
		} else fprintf(out, "NA,NA,NA,");

		// CHISQ and P-value 
		if (pvalue >= 0 && pvalue <= 1.0) 
			fprintf(out, "%f,%f,", chisq, pvalue);
		else
			fprintf(out, "NA,NA,");

		// EMP1 & EMP2
		fprintf(out, "%10.16f,%10.16f,%f,%f,%f\n", emp1, emp2, adj->bonf, adj->holm, adj->fdr_bh);
	}

	return chisq;
}
